import os
import pandas as pd
import numpy as np
import torch
from time import time
import warnings

# 忽略警告
warnings.filterwarnings("ignore")

# 训练数据的拼接函数
def collate_fn_train(data):
    # 按照历史点击序列的长度降序排序
    data.sort(key=lambda x: len(x[1]), reverse=True)
    # 计算每个序列的长度
    lens = [len(hist) for _, hist, _ in data]
    users, labels = [], []
    # 初始化一个全零的张量，用于存储填充后的序列
    padded_seq = torch.zeros(len(data), max(lens)).long()
    for i, (user, hist, label) in enumerate(data):
        users.append(user)
        # 将历史点击序列填充到对应位置
        padded_seq[i, :lens[i]] = torch.LongTensor(hist)
        labels.append(label)

    return users, padded_seq, torch.tensor(labels).long(), lens

# 验证数据的拼接函数
def collate_fn_vali(data):
    # 按照历史点击序列的长度降序排序
    data.sort(key=lambda x: len(x[1]), reverse=True)
    # 计算每个序列的长度
    lens = [len(hist) for _, hist, _, _ in data]
    users, labels = [], []
    # 初始化一个全零的张量，用于存储填充后的序列
    padded_seq = torch.zeros(len(data), max(lens)).long()
    click_mask_matrix = []
    for i, (user, hist, label, cmm) in enumerate(data):
        users.append(user)
        # 将历史点击序列填充到对应位置
        padded_seq[i, :lens[i]] = torch.LongTensor(hist)
        labels.append(label)
        # 将点击掩码矩阵添加到列表
        click_mask_matrix.append(cmm.reshape(1, -1))
    # 将所有的点击掩码矩阵按行拼接
    click_mask_matrix = torch.cat(click_mask_matrix, dim=0)

    return users, padded_seq, torch.tensor(labels).long(), lens, click_mask_matrix

# 测试数据的拼接函数
def collate_fn_test(data):
    # 按照历史点击序列的长度降序排序
    data.sort(key=lambda x: len(x[1]), reverse=True)
    # 计算每个序列的长度
    lens = [len(hist) for _, hist, _ in data]
    users = []
    # 初始化一个全零的张量，用于存储填充后的序列
    padded_seq = torch.zeros(len(data), max(lens)).long()
    click_mask_matrix = []
    for i, (user, hist, cmm) in enumerate(data):
        users.append(user)
        # 将历史点击序列填充到对应位置
        padded_seq[i, :lens[i]] = torch.LongTensor(hist)
        # 将点击掩码矩阵添加到列表
        click_mask_matrix.append(cmm.reshape(1, -1))
    # 将所有的点击掩码矩阵按行拼接
    click_mask_matrix = torch.cat(click_mask_matrix, dim=0)

    return users, padded_seq, lens, click_mask_matrix
